module latlon_data_mod

  use datetime
  use const_mod
  use latlon_process_mod
  use latlon_parallel_mod

  implicit none

  private

  public latlon_data_type

  type latlon_data_type
    integer :: eid = 0 ! Ensemble ID
    type(datetime_type) time
    ! Surface variables
    real(r8), allocatable, dimension(:,:  ) :: zs
    real(r8), allocatable, dimension(:,:  ) :: ps
    real(r8), allocatable, dimension(:,:  ) :: t2
    real(r8), allocatable, dimension(:,:  ) :: q2
    real(r8), allocatable, dimension(:,:  ) :: rh2
    real(r8), allocatable, dimension(:,:  ) :: u10
    real(r8), allocatable, dimension(:,:  ) :: v10
    ! 3D variables on isobaric surfaces
    real(r8), allocatable, dimension(:,:,:) :: u
    real(r8), allocatable, dimension(:,:,:) :: v
    real(r8), allocatable, dimension(:,:,:) :: t
    real(r8), allocatable, dimension(:,:,:) :: q
    real(r8), allocatable, dimension(:,:,:) :: z
  contains
    procedure :: init => latlon_data_init
    procedure :: clear => latlon_data_clear
    procedure :: write => latlon_data_write
    procedure, private :: latlon_data_assign
    generic :: assignment(=) => latlon_data_assign
    procedure :: zero => latlon_data_zero
    procedure :: add => latlon_data_add
    procedure :: sub => latlon_data_sub
    procedure :: divide => latlon_data_divide
    procedure :: global_sum => latlon_global_sum
    final :: latlon_data_final
  end type latlon_data_type

contains

  subroutine latlon_data_init(this, eid)

    class(latlon_data_type), intent(inout) :: this
    integer, intent(in) :: eid

    integer lon_lb, lon_ub, lat_lb, lat_ub, lev_lb, lev_ub

    call this%clear()

    this%eid = eid

    lon_lb = proc%local_mesh%lon_lb
    lon_ub = proc%local_mesh%lon_ub
    lat_lb = proc%local_mesh%lat_lb
    lat_ub = proc%local_mesh%lat_ub
    lev_lb = proc%local_mesh%lev_lb
    lev_ub = proc%local_mesh%lev_ub

    allocate(this%zs (lon_lb:lon_ub,lat_lb:lat_ub))
    allocate(this%ps (lon_lb:lon_ub,lat_lb:lat_ub))
    allocate(this%t2 (lon_lb:lon_ub,lat_lb:lat_ub))
    allocate(this%q2 (lon_lb:lon_ub,lat_lb:lat_ub))
    allocate(this%rh2(lon_lb:lon_ub,lat_lb:lat_ub))
    allocate(this%u10(lon_lb:lon_ub,lat_lb:lat_ub))
    allocate(this%v10(lon_lb:lon_ub,lat_lb:lat_ub))
    allocate(this%u  (lon_lb:lon_ub,lat_lb:lat_ub,lev_lb:lev_ub))
    allocate(this%v  (lon_lb:lon_ub,lat_lb:lat_ub,lev_lb:lev_ub))
    allocate(this%t  (lon_lb:lon_ub,lat_lb:lat_ub,lev_lb:lev_ub))
    allocate(this%q  (lon_lb:lon_ub,lat_lb:lat_ub,lev_lb:lev_ub))
    allocate(this%z  (lon_lb:lon_ub,lat_lb:lat_ub,lev_lb:lev_ub))

  end subroutine latlon_data_init

  subroutine latlon_data_clear(this)

    class(latlon_data_type), intent(inout) :: this

    if (allocated(this%zs )) deallocate(this%zs )
    if (allocated(this%ps )) deallocate(this%ps )
    if (allocated(this%t2 )) deallocate(this%t2 )
    if (allocated(this%q2 )) deallocate(this%q2 )
    if (allocated(this%rh2)) deallocate(this%rh2)
    if (allocated(this%u10)) deallocate(this%u10)
    if (allocated(this%v10)) deallocate(this%v10)
    if (allocated(this%u  )) deallocate(this%u  )
    if (allocated(this%v  )) deallocate(this%v  )
    if (allocated(this%t  )) deallocate(this%t  )
    if (allocated(this%q  )) deallocate(this%q  )
    if (allocated(this%z  )) deallocate(this%z  )

  end subroutine latlon_data_clear

  subroutine latlon_data_write(this, file_path)

    use fiona
    use string

    class(latlon_data_type), intent(in) :: this
    character(*), intent(in) :: file_path

    character(8) key
    integer is, ie, js, je, ks, ke
    integer start(3), count(3)

    key = 'ens' // to_str(this%eid)

    call fiona_create_dataset(key, file_path=file_path, start_time=this%time%isoformat(), time_units='hours', mpi_comm=proc%comm)
    call fiona_add_dim(key, 'time', add_var=.true.)
    call fiona_add_dim(key, 'lon' , size=proc%global_mesh%num_lon, add_var=.true., decomp=.true.)
    call fiona_add_dim(key, 'lat' , size=proc%global_mesh%num_lat, add_var=.true., decomp=.true.)
    call fiona_add_dim(key, 'lev' , size=proc%global_mesh%num_lev, add_var=.true.)
    call fiona_add_var(key, 'ps'  , long_name='surface pressure'        , units='hPa'   , dim_names=['lon ', 'lat ',         'time'])
    call fiona_add_var(key, 't2'  , long_name='temperature at 2m'       , units='K'     , dim_names=['lon ', 'lat ',         'time'])
    call fiona_add_var(key, 'q2'  , long_name='specific humidity at 2m' , units='Kg Kg-1', dim_names=['lon ', 'lat ',         'time'])
    call fiona_add_var(key, 'rh2' , long_name='relative humidity at 2m' , units='%'     , dim_names=['lon ', 'lat ',         'time'])
    call fiona_add_var(key, 'u10' , long_name='u wind component at 10m' , units='m s-1' , dim_names=['lon ', 'lat ',         'time'])
    call fiona_add_var(key, 'v10' , long_name='v wind component at 10m' , units='m s-1' , dim_names=['lon ', 'lat ',         'time'])
    call fiona_add_var(key, 'u'   , long_name='u wind component'        , units='m s-1' , dim_names=['lon ', 'lat ', 'lev ', 'time'])
    call fiona_add_var(key, 'v'   , long_name='v wind component'        , units='m s-1' , dim_names=['lon ', 'lat ', 'lev ', 'time'])
    call fiona_add_var(key, 't'   , long_name='temperature'             , units='K'     , dim_names=['lon ', 'lat ', 'lev ', 'time'])
    call fiona_add_var(key, 'q'   , long_name='specific humidity'       , units='Kg Kg-1', dim_names=['lon ', 'lat ', 'lev ', 'time'])
    call fiona_start_output(key, 0.0d0, new_file=.true.)
    call fiona_output(key, 'lon'  , proc%global_mesh%lon_deg(1:proc%global_mesh%num_lon))
    call fiona_output(key, 'lat'  , proc%global_mesh%lat_deg(1:proc%global_mesh%num_lat))
    call fiona_output(key, 'lev'  , proc%global_mesh%lev    (1:proc%global_mesh%num_lev))

    is = proc%local_mesh%lon_ibeg
    ie = proc%local_mesh%lon_iend
    js = proc%local_mesh%lat_ibeg
    je = proc%local_mesh%lat_iend
    ks = proc%local_mesh%lev_ibeg
    ke = proc%local_mesh%lev_iend
    start = [is,js,ks]
    count = [proc%local_mesh%num_lon,proc%local_mesh%num_lat,proc%local_mesh%num_lev]

    call fiona_output(key, 'ps'   , this%ps (is:ie,js:je      ), start=start, count=count)
    call fiona_output(key, 't2'   , this%t2 (is:ie,js:je      ), start=start, count=count)
    call fiona_output(key, 'q2'   , this%q2 (is:ie,js:je      ), start=start, count=count)
    call fiona_output(key, 'rh2'  , this%rh2(is:ie,js:je      ), start=start, count=count)
    call fiona_output(key, 'u10'  , this%u10(is:ie,js:je      ), start=start, count=count)
    call fiona_output(key, 'v10'  , this%v10(is:ie,js:je      ), start=start, count=count)
    call fiona_output(key, 'u'    , this%u  (is:ie,js:je,ks:ke), start=start, count=count)
    call fiona_output(key, 'v'    , this%v  (is:ie,js:je,ks:ke), start=start, count=count)
    call fiona_output(key, 't'    , this%t  (is:ie,js:je,ks:ke), start=start, count=count)
    call fiona_output(key, 'q'    , this%q  (is:ie,js:je,ks:ke), start=start, count=count)
    call fiona_end_output(key)

  end subroutine latlon_data_write

  subroutine latlon_data_assign(this, other)

    class(latlon_data_type), intent(inout) :: this
    class(latlon_data_type), intent(in) :: other

    this%ps = other%ps
    this%t2 = other%t2
    this%q2 = other%q2
    this%rh2 = other%rh2
    this%u10 = other%u10
    this%v10 = other%v10
    this%u = other%u
    this%v = other%v
    this%t = other%t
    this%q = other%q
    this%z = other%z

  end subroutine latlon_data_assign

  subroutine latlon_data_zero(this)

    class(latlon_data_type), intent(inout) :: this

    this%ps = 0
    this%t2 = 0
    this%q2 = 0
    this%rh2 = 0
    this%u10 = 0
    this%v10 = 0
    this%u = 0
    this%v = 0
    this%t = 0
    this%q = 0
    this%z = 0

  end subroutine latlon_data_zero

  subroutine latlon_data_add(this, other)

    class(latlon_data_type), intent(inout) :: this
    type(latlon_data_type), intent(in) :: other

    this%ps = this%ps + other%ps
    this%t2 = this%t2 + other%t2
    this%q2 = this%q2 + other%q2
    this%rh2 = this%rh2 + other%rh2
    this%u10 = this%u10 + other%u10
    this%v10 = this%v10 + other%v10
    this%u = this%u + other%u
    this%v = this%v + other%v
    this%t = this%t + other%t
    this%q = this%q + other%q
    this%z = this%z + other%z

  end subroutine latlon_data_add

  subroutine latlon_data_sub(this, other)

    class(latlon_data_type), intent(inout) :: this
    type(latlon_data_type), intent(in) :: other

    this%ps = this%ps - other%ps
    this%t2 = this%t2 - other%t2
    this%q2 = this%q2 - other%q2
    this%rh2 = this%rh2 - other%rh2
    this%u10 = this%u10 - other%u10
    this%v10 = this%v10 - other%v10
    this%u = this%u - other%u
    this%v = this%v - other%v
    this%t = this%t - other%t
    this%q = this%q - other%q
    this%z = this%z - other%z

  end subroutine latlon_data_sub

  subroutine latlon_data_divide(this, scalar)

    class(latlon_data_type), intent(inout) :: this
    real(r8), intent(in) :: scalar

    this%ps = this%ps / scalar
    this%t2 = this%t2 / scalar
    this%q2 = this%q2 / scalar
    this%rh2 = this%rh2 / scalar
    this%u10 = this%u10 / scalar
    this%v10 = this%v10 / scalar
    this%u = this%u / scalar
    this%v = this%v / scalar
    this%t = this%t / scalar
    this%q = this%q / scalar
    this%z = this%z / scalar

  end subroutine latlon_data_divide

  subroutine latlon_global_sum(this)

    class(latlon_data_type), intent(inout) :: this

    call global_sum(proc%comm, this%ps )
    call global_sum(proc%comm, this%t2 )
    call global_sum(proc%comm, this%q2 )
    call global_sum(proc%comm, this%rh2)
    call global_sum(proc%comm, this%u10)
    call global_sum(proc%comm, this%v10)
    call global_sum(proc%comm, this%u  )
    call global_sum(proc%comm, this%v  )
    call global_sum(proc%comm, this%t  )
    call global_sum(proc%comm, this%q  )
    call global_sum(proc%comm, this%z  )

  end subroutine latlon_global_sum

  subroutine latlon_data_final(this)

    type(latlon_data_type), intent(inout) :: this

    call this%clear()

  end subroutine latlon_data_final

end module latlon_data_mod
